# coding: utf-8
from gensim import corpora
import matplotlib.pyplot as plt
import matplotlib
import warnings
import pyLDAvis.gensim_models

warnings.filterwarnings('ignore')

from gensim.models.coherencemodel import CoherenceModel
from gensim.models.ldamodel import LdaModel
from nltk.corpus import stopwords

def load_data():
    file = open("OpenEuler_description.csv", encoding="utf-8", errors='ignore').read().split("\n")
    data_set = []
    for i in range(len(file)):
        result = []
        seg_list = file[i].split()
        for w in seg_list:
            result.append(w)
        filtered_words = [word for word in result if word not in stopwords.words('english')]
        data_set.append(filtered_words)
    dictionary = corpora.Dictionary(data_set)
    corpus = [dictionary.doc2bow(text) for text in data_set]# 表示为第几个单词出现了几次
    return data_set, dictionary, corpus


def coherence(num_topics):
    data_set, dictionary, corpus = load_data()
    ldamodel = LdaModel(corpus, num_topics=num_topics, id2word=dictionary, passes=30, random_state=1)
    ldacm = CoherenceModel(model=ldamodel, texts=data_set, dictionary=dictionary, coherence='c_v')
    print(ldacm.get_coherence())
    return ldacm.get_coherence()

def draw():
    x = range(1,15)
    y = [coherence(i) for i in x]
    plt.plot(x, y)
    plt.xlabel("主题数目")
    plt.ylabel('coherence大小')
    plt.rcParams['font.sans-serif'] = ['SimHei']
    matplotlib.rcParams['axes.unicode_minus'] = False
    plt.title('主题-coherence变化情况')
    plt.show()

def build():
    data_set, dictionary, corpus = load_data()
    lda = LdaModel(corpus=corpus, id2word=dictionary, num_topics=9, passes=30, random_state=1)
    topic_list = lda.print_topics()
    print(topic_list)

    for i in lda.get_document_topics(corpus)[:]:
        listj = []
        for j in i:
            listj.append(j[1])
        bz = listj.index(max(listj))
        print(i[bz][0])
    #pyLDAvis.enable_notebook()
    data = pyLDAvis.gensim_models.prepare(lda, corpus, dictionary, mds='mmds')
    pyLDAvis.save_html(data, '9topic.html')


if __name__ == '__main__':
    build()
